# 1.图像预测
import argparse
import os
import threading
import time

from imageai.Classification import ImageClassification


def _argparse():
    parse = argparse.ArgumentParser()
    parse.add_argument("--model_name", default='resnet50-19c8e357.pth', action='store', required=False,
                       dest="model_name",
                       help="模型，默认resnet50-19c8e357.pth")
    parse.add_argument("--type", default='single', action='store', required=False, dest="type",
                       help="模式，默认识别一个图片")
    parse.add_argument("--file_path", default='ImagePrediction\\1.jpg', action='store', required=False,
                       dest="file_path",
                       help="图片位置，默认type=single为图片地址，type=multiThreading时为图片所在目录")
    return parse.parse_args()


def main():
    parser = _argparse()

    prediction = ImageClassification()
    if parser.model_name == 'mobilenet_v2-b0353104.pth':
        prediction.setModelTypeAsMobileNetV2()
    elif parser.model_name == 'inception_v3_google-1a9a5a14.pth':
        prediction.setModelTypeAsInceptionV3()
    elif parser.model_name == 'densenet121-a639ec97.pth':
        prediction.setModelTypeAsDenseNet121()
    else:
        prediction.setModelTypeAsResNet50()

    models = "models"
    execution_path = os.getcwd()
    model_path = os.path.join(execution_path, models, parser.model_name)

    prediction.setModelPath(model_path)
    prediction.loadModel()

    if parser.type == 'single':
        print("开始时间为:", time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
        predictions, probabilities = prediction.classifyImage(os.path.join(execution_path, parser.file_path),
                                                              result_count=5)
        print("结束时间为:", time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))

        for eachPrediction, eachProbability in zip(predictions, probabilities):
            print(eachPrediction, " : ", eachProbability)
    elif parser.type == "multiThreading":
        picturesfolder = os.path.join(execution_path, parser.file_path)
        allfiles = os.listdir(picturesfolder)

        class PredictionThread(threading.Thread):
            def init(self):
                threading.Thread.init(self)

            def run(self):
                for eachPicture in allfiles:
                    if eachPicture.endswith(".png") or eachPicture.endswith(".jpg"):
                        predictions, probabilities = prediction.classifyImage(os.path.join(picturesfolder, eachPicture),
                                                                              result_count=1)
                        for eachPrediction, eachProbability in zip(predictions, probabilities):
                            print(eachPicture, " : ", eachPrediction, " : ", eachProbability)

        pt = PredictionThread()
        pt.start()
    else:
        print("params type out of range")


if __name__ == '__main__':
    main()
